AWS Advanced NodeJS WrapperのDSQL用プラグインを実装してみた

岩田智哉

2024.12.27

リテールアプリ共創部@大阪の岩田です。AWS Advanced NodeJS Wrapperのプラグインを実装してDSQLに接続してみたので、その手順をご紹介します。

環境

今回利用した環境は以下の通りです。

Node.js v22.11.0
@aws-sdk/dsql-signer: 3.716.0
aws-advanced-nodejs-wrapper: 1.1.0
pg: 8.13.1

やってみる

AWS Advanced NodeJS Wrapperのプラグインを実装する手順については以下のドキュメントで紹介されています。

この手順に従いつつ、IamAuthenticationPluginの実装を参考にすることで簡易なDSQL用プラグインが実装できました。以後順を追いながら作業していきます。

依存ライブラリのインストール

まずは必要なライブラリ類をインストールします。

npm install @aws-sdk/dsql-signer \
  aws-advanced-nodejs-wrapper \
  pg

aws-advanced-nodejs-wrapper とpgに加えて@aws-sdk/dsql-signerをインストールしています。このライブラリを利用するとDSQLに接続するための一時トークンを簡単に生成可能です。利用手順については以下のドキュメントが参考になります。

PluginFactoryの実装

まずDSQL用プラグインのインスタンスを生成するためにConnectionPluginFactoryを継承したクラスを定義します。このクラスはgetInstanceというメソッドでプラグインのインスタンスを生成してreturnする必要があります。今回は単にプラグインのインスタンスをnew()して返却する実装としました。

export class DsqlAuthenticationPluginFactory extends ConnectionPluginFactory {
  async getInstance(
    pluginService: PluginService,
    properties: object,
  ): Promise<ConnectionPlugin> {
    return new DsqlAuthenticationPlugin(pluginService);
  }
}

new()しているDsqlAuthenticationPluginクラスは後ほど実装します。

プラグインの登録

続いて先ほど実装したDsqlAuthenticationPluginFactoryをAWS Advanced NodeJS Wrapperのプラグインとして登録します。PluginManager.registerPluginの第1引数にプラグインの名前を、第2引数にプラグインのインスタンスを生成するためのFactoryクラスを渡して呼び出します。

PluginManager.registerPlugin("dsql", DsqlAuthenticationPluginFactory);

これでdsqlという名前で自作プラグインが利用可能になりました。次にクライアントクラスを生成します。

const postgresHost = "<DSQLのエンドポイント>";
const client = new AwsPGClient({
  host: postgresHost,
  port: 5432,
  database: "postgres",
  user: "admin",
  plugins: "dsql",
  ssl: true,
});

plugins: "dsql"の指定によってこのクライアントクラスのインスタンスが自作DSQL用プラグインを利用してくれるようになります。

DsqlAuthenticationPluginの実装

続いてメイン部分であるDsqlAuthenticationPluginクラスを実装します。このクラスはDBとの接続に利用するためのプラグインなのでAbstractConnectionPluginを継承します。

class DsqlAuthenticationPlugin extends AbstractConnectionPlugin {
  private static readonly SUBSCRIBED_METHODS = new Set<string>([
    "connect",
    "forceConnect",
  ]);
  private pluginService: PluginService;

  constructor(pluginService: PluginService) {
    super();
    this.pluginService = pluginService;
  }

  getSubscribedMethods(): Set<string> {
    return DsqlAuthenticationPlugin.SUBSCRIBED_METHODS;
  }
}

AbstractConnectionPluginはgetSubscribedMethodsというメソッドが継承必須となっているため、まずはこのメソッドを実装します。戻り値はプラグインに実装するメソッドの一覧をstringのSetとして返却します。ここで定義したメソッドがPluginManagerによって呼び出されるので、自前のロジックをAwsPGClientの各種処理に適宜差し込めます。この辺りのアーキテクチャは以下のドキュメントが参考になります。

今回はconnectとforceConnectを指定したので、これらのメソッドも実装します。

class DsqlAuthenticationPlugin extends AbstractConnectionPlugin {
//...略  
	connect(
    hostInfo: HostInfo,
    props: Map<string, any>,
    isInitialConnection: boolean,
    connectFunc: () => Promise<ClientWrapper>,
  ): Promise<ClientWrapper> {
    return this.connectInternal(
      hostInfo,
      props,
      isInitialConnection,
      connectFunc,
    );
  }

  forceConnect(
    hostInfo: HostInfo,
    props: Map<string, any>,
    isInitialConnection: boolean,
    forceConnectFunc: () => Promise<ClientWrapper>,
  ): Promise<ClientWrapper> {
    return this.connectInternal(
      hostInfo,
      props,
      isInitialConnection,
      forceConnectFunc,
    );
  }
//...略
}

connectとforceConnectともにプライベートなconnectInternalというメソッドを呼びだすだけの実装です。実際のロジックはconnectInternalに実装します。

メイン処理のconnectInternalです。

class DsqlAuthenticationPlugin extends AbstractConnectionPlugin {
//...略  
	private async connectInternal(
    hostInfo: HostInfo,
    props: Map<string, any>,
    isInitialConnection: boolean,
    connectFunc: () => Promise<ClientWrapper>,
  ): Promise<ClientWrapper> {
    const region = hostInfo.host.split(".")[2];
    const signer = new DsqlSigner({
      hostname: hostInfo.host,
      region,
    });
    const token = await signer.getDbConnectAdminAuthToken();
    WrapperProperties.PASSWORD.set(props, token);
    this.pluginService.updateConfigWithProperties(props);
    return connectFunc();
  }
}

引数のhostInfoで接続先ホストの情報が渡されてくるのでhostInfo.host.split(".")[2];でリージョンを取り出し、DsqlSignerクラスを使って一時トークンを生成します。生成した一時トークンをPluginServiceクラスのupdateConfigWithPropertiesでパスワードとしてセットし、connectFunc();でオリジナルのDB接続処理を呼び出します。

参考にしたIamAuthenticationPluginでは一時トークンのキャッシュなど高度なことを色々とやっているのですが、今回はプラグイン実装について理解を深めるためのサンプル実装なので愚直に接続の都度一時トークンを生成する実装としています。

最終形

最終的なコードは以下のようになりました。

import { AwsPGClient } from "aws-advanced-nodejs-wrapper/dist/pg/lib/index.js";
import {
  PluginManager,
  ConnectionPlugin,
} from "aws-advanced-nodejs-wrapper/dist/common/lib/index.js";
import { ConnectionPluginFactory } from "aws-advanced-nodejs-wrapper/dist/common/lib/plugin_factory";
import { AbstractConnectionPlugin } from "aws-advanced-nodejs-wrapper/dist/common/lib/abstract_connection_plugin";
import { HostInfo } from "aws-advanced-nodejs-wrapper/dist/common/lib/host_info";
import { ClientWrapper } from "aws-advanced-nodejs-wrapper/dist/common/lib/client_wrapper";
import { WrapperProperties } from "aws-advanced-nodejs-wrapper/dist/common/lib/wrapper_property";
import { PluginService } from "aws-advanced-nodejs-wrapper/dist/common/lib/plugin_service";
import { DsqlSigner } from "@aws-sdk/dsql-signer";

class DsqlAuthenticationPlugin extends AbstractConnectionPlugin {
  private static readonly SUBSCRIBED_METHODS = new Set<string>([
    "connect",
    "forceConnect",
  ]);
  private pluginService: PluginService;

  constructor(pluginService: PluginService) {
    super();
    this.pluginService = pluginService;
  }

  getSubscribedMethods(): Set<string> {
    return DsqlAuthenticationPlugin.SUBSCRIBED_METHODS;
  }

  connect(
    hostInfo: HostInfo,
    props: Map<string, any>,
    isInitialConnection: boolean,
    connectFunc: () => Promise<ClientWrapper>,
  ): Promise<ClientWrapper> {
    return this.connectInternal(
      hostInfo,
      props,
      isInitialConnection,
      connectFunc,
    );
  }

  forceConnect(
    hostInfo: HostInfo,
    props: Map<string, any>,
    isInitialConnection: boolean,
    forceConnectFunc: () => Promise<ClientWrapper>,
  ): Promise<ClientWrapper> {
    return this.connectInternal(
      hostInfo,
      props,
      isInitialConnection,
      forceConnectFunc,
    );
  }

  private async connectInternal(
    hostInfo: HostInfo,
    props: Map<string, any>,
    isInitialConnection: boolean,
    connectFunc: () => Promise<ClientWrapper>,
  ): Promise<ClientWrapper> {
    const region = hostInfo.host.split(".")[2];
    const signer = new DsqlSigner({
      hostname: hostInfo.host,
      region,
    });
    const token = await signer.getDbConnectAdminAuthToken();
    WrapperProperties.PASSWORD.set(props, token);
    this.pluginService.updateConfigWithProperties(props);
    return connectFunc();
  }
}

export class DsqlAuthenticationPluginFactory extends ConnectionPluginFactory {
  async getInstance(
    pluginService: PluginService,
    properties: object,
  ): Promise<ConnectionPlugin> {
    return new DsqlAuthenticationPlugin(pluginService);
  }
}

PluginManager.registerPlugin("dsql", DsqlAuthenticationPluginFactory);

const main = async () => {
  const postgresHost = "<DSQLのエンドポイント>";
  const client = new AwsPGClient({
    host: postgresHost,
    port: 5432,
    database: "postgres",
    user: "admin",
    plugins: "dsql",
    ssl: true,
  });

  try {
    await client.connect();
    const result = await client.query("select now()");
    console.log(result);
  } finally {
    await client.end();
  }
};

main();

動作確認

実装できたので動作確認してみます。今回TypeScriptで実装したのでtsxで実行しました。

❯ npm start

> start
> tsx main.ts

結果は以下の通りでした。無事にDSQLに接続してSQLが発行できています！

Result {
  command: 'SELECT',
  rowCount: 1,
  oid: null,
  rows: [ { now: 2024-12-27T02:36:03.829Z } ],
  fields: [
    Field {
      name: 'now',
      tableID: 0,
      columnID: 0,
      dataTypeID: 1184,
      dataTypeSize: 8,
      dataTypeModifier: -1,
      format: 'text'
    }
  ],
  _parsers: [ [Function: parseDate] ],
  _types: TypeOverrides {
    _types: {
      getTypeParser: [Function: getTypeParser],
      setTypeParser: [Function: setTypeParser],
      arrayParser: [Object],
      builtins: [Object]
    },
    text: {},
    binary: {}
  },
  RowCtor: null,
  rowAsArray: false,
  _prebuiltEmptyResultObject: { now: null }
}

まとめ

AWS Advanced NodeJS WrapperのDSQL用プラグインを実装してみました。

まあ、しばらく待てば公式からリリースされそうな気はしますが、今すぐAWS Advanced NodeJS WrapperからDSQLに接続したいという方には参考になるかもしれません。もし実際に利用する場合はIamAuthenticationPluginを参考にキャッシュロジックなどをしっかり作りこんでから利用するようにお願いします。